package main

import (
	"bufio"
	"database/sql"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"log"
	"os"
	"strconv"
	"sync"
	"time"
)

//链接数据池
//先根据条件查出所有的工单条数,total
//然后根据工单条数还有线程数进行分割,平均分配给每个线程,线程数thread_num
//每个线程执行的数据量avg_num = total/thread_num
//每个线程每次执行的数量batch_size
//先创建goroutine数量为thread_num
//每个携程中执行导出单个csv命令,其中每条导出命令是分页的,开始页码也就是
//pageNo是根据他处于的携程数来计算的,比如他处于第三个携程,那么他最开始的pageNo = (3-1)*avg_num,
//有一点需要注意的是如果该携程处于最后一个携程那么他的pageSize可能不是batch_size
const (
	OUTPUT_DIR = "./output/" // 输出文件目录
	THREAD_NUM = 12          // 线程数
	BATCH_SIZE = 1000000     // 每次查询的记录数
)

func main() {
	start := time.Now()
	//链接数据库
	db, err := sql.Open("mysql", "root:123456@tcp(localhost:3306)/my_test")
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()
	//获取数据总量
	var total int
	if err := db.QueryRow("SELECT COUNT(*) FROM users").Scan(&total); err != nil {
		log.Fatal(err)
	}
	//计算每个携程处理的范围
	blockSize := total / THREAD_NUM
	blocks := make([][2]int, THREAD_NUM)
	for b := 0; b < THREAD_NUM; b++ {
		blocks[b][0] = b * blockSize
		blocks[b][1] = b*blockSize + blockSize
		if b == THREAD_NUM {
			blocks[b][1] = total
		}
	}
	var wg sync.WaitGroup
	//开启携程
	for i := 0; i < THREAD_NUM; i++ {
		wg.Add(1)
		go func(block [2]int, filename string) {
			defer wg.Done()

			//开始导出csv
			file, err := os.Create(filename)
			if err != nil {
				log.Fatal(err)
			}
			defer file.Close()

			writer := bufio.NewWriter(file)
			defer writer.Flush()

			// 写入 CSV 文件头
			headers := []string{"id", "name", "age", "email"}
			writer.WriteString(fmt.Sprintf("%s\n", headers))

			for start := block[0]; start < block[1]; start += BATCH_SIZE {
				end := start + BATCH_SIZE
				if end > block[1] {
					end = block[1]
				}
				rows, err := db.Query("SELECT id,name,age,email FROM users LIMIT ?, ?", start, end-start)
				if err != nil {
					log.Fatal(err)
				}
				// 处理查询结果
				for rows.Next() {
					var id int
					var name string
					var age int
					var email string
					if err := rows.Scan(&id, &name, &age, &email); err != nil {
						log.Fatal(err)
					}

					// 将记录写入小文件
					record := []string{strconv.Itoa(id), name, strconv.Itoa(age), email}
					writer.WriteString(fmt.Sprintf("%s\n", record))
				}
				if err := rows.Err(); err != nil {
					log.Fatal(err)
				}
			}
		}(blocks[i], fmt.Sprintf("%s%d.csv", OUTPUT_DIR, i))
	}
	wg.Wait()

	//合并到一个大文件
	outputFile := fmt.Sprintf("%s%s.csv", OUTPUT_DIR, "output")
	outFile, err := os.Create(outputFile)
	if err != nil {
		log.Fatal(err)
	}
	defer outFile.Close()

	writer := bufio.NewWriter(outFile)
	defer writer.Flush()

	for i := 0; i < THREAD_NUM; i++ {
		//按照顺序依次打开小的csv
		file, err := os.Open(fmt.Sprintf("%s%d.csv", OUTPUT_DIR, i))
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()

		//用扫描器开始读取文件
		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			record := scanner.Text()
			if len(record) == 0 {
				continue
			}
			writer.WriteString(fmt.Sprintf("%s\n", record))
		}
	}
	fmt.Printf("处理完成，结果保存在文件 %s 中\n", outputFile)
	end := time.Now()
	fmt.Printf("Time elapsed: %v\n", end.Sub(start))
}
